{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2.6 Stokes equation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Find $u \\in [H^1_D]^2$ and $p \\in L_2$ such that\n",
    "\n",
    "$$\n",
    "\\DeclareMathOperator{\\Div}{div}\n",
    "\\begin{array}{ccccll}\n",
    "\\int \\nabla u : \\nabla v & + & \\int \\Div v \\, p & = & \\int f v & \\forall \\, v \\\\\n",
    "\\int \\Div u \\, q &  &  & = & 0 & \\forall \\, q\n",
    "\\end{array}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define channel geometry and mesh it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ngsolve import *\n",
    "import netgen.gui\n",
    "%gui tk\n",
    "\n",
    "from netgen.geom2d import SplineGeometry\n",
    "geo = SplineGeometry()\n",
    "geo.AddRectangle( (0, 0), (2, 0.41), bcs = (\"wall\", \"outlet\", \"wall\", \"inlet\"))\n",
    "geo.AddCircle ( (0.2, 0.2), r=0.05, leftdomain=0, rightdomain=1, bc=\"cyl\")\n",
    "mesh = Mesh( geo.GenerateMesh(maxh=0.05))\n",
    "mesh.Curve(3)\n",
    "Draw (mesh)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use Taylor Hood finite element pairing: Continuous $P^2$ elements for velocity, and continuous $P^1$ for pressure:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "V = H1(mesh, order=2, dirichlet=\"wall|inlet|cyl\")\n",
    "Q = H1(mesh, order=1)\n",
    "X = FESpace([V,V,Q])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Setup bilinear-form for Stokes. We give names for all scalar field components. The divergence is constructed from partial derivatives of the velocity components."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "ux,uy,p = X.TrialFunction()\n",
    "vx,vy,q = X.TestFunction()\n",
    "\n",
    "div_u = grad(ux)[0]+grad(uy)[1]\n",
    "div_v = grad(vx)[0]+grad(vy)[1]\n",
    "\n",
    "a = BilinearForm(X)\n",
    "a += SymbolicBFI(grad(ux)*grad(vx)+grad(uy)*grad(vy) + div_u*q + div_v*p)\n",
    "a.Assemble()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set inhomogeneous Dirichlet boundary condition only on inlet boundary:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "gfu = GridFunction(X)\n",
    "uin = 1.5*4*y*(0.41-y)/(0.41*0.41)\n",
    "gfu.components[0].Set(uin, definedon=mesh.Boundaries(\"inlet\"))\n",
    "velocity = CoefficientFunction(gfu.components[0:2])\n",
    "Draw(velocity, mesh, \"vel\")\n",
    "Draw(Norm(velocity), mesh, \"|vel|\")\n",
    "SetVisualization(max=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solve equation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "res = gfu.vec.CreateVector()\n",
    "res.data = -a.mat * gfu.vec\n",
    "inv = a.mat.Inverse(freedofs=X.FreeDofs(), inverse=\"umfpack\")\n",
    "gfu.vec.data += inv * res\n",
    "Redraw()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Testing different velocity-pressure pairs\n",
    "Now we define a Stokes setup function to test different spaces:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def SolveStokes(X):\n",
    "    ux,uy,p = X.TrialFunction()\n",
    "    vx,vy,q = X.TestFunction()\n",
    "    div_u = grad(ux)[0]+grad(uy)[1]\n",
    "    div_v = grad(vx)[0]+grad(vy)[1]\n",
    "    a = BilinearForm(X)\n",
    "    a += SymbolicBFI(grad(ux)*grad(vx)+grad(uy)*grad(vy) + div_u*q + div_v*p)\n",
    "    a.Assemble()\n",
    "    gfu = GridFunction(X)\n",
    "    uin = 1.5*4*y*(0.41-y)/(0.41*0.41)\n",
    "    gfu.components[0].Set(uin, definedon=mesh.Boundaries(\"inlet\"))\n",
    "    res = gfu.vec.CreateVector()\n",
    "    res.data = -a.mat * gfu.vec\n",
    "    inv = a.mat.Inverse(freedofs=X.FreeDofs(), inverse=\"umfpack\")\n",
    "    gfu.vec.data += inv * res\n",
    "    \n",
    "\n",
    "    velocity = CoefficientFunction(gfu.components[0:2])\n",
    "    Draw(velocity, mesh, \"vel\")\n",
    "    Draw(Norm(velocity), mesh, \"|vel|\")\n",
    "    SetVisualization(max=2)    \n",
    "    \n",
    "    return gfu"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Higher order Taylor-Hood elements:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "V = H1(mesh, order=4, dirichlet=\"wall|inlet|cyl\")\n",
    "Q = H1(mesh, order=3)\n",
    "X = FESpace([V,V,Q])\n",
    "\n",
    "gfu = SolveStokes(X)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With discontinuous pressure elements P2-P1 is unstable:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "V.ndof = 1702 , Q.ndof = 2382\n"
     ]
    },
    {
     "ename": "RuntimeError",
     "evalue": "UmfpackInverse: Numeric factorization failed.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-8-4e71f02dc134>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0mX\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mFESpace\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mV\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mV\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mQ\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 6\u001b[0;31m \u001b[0mgfu\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mSolveStokes\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mX\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m<ipython-input-6-13793926ffa5>\u001b[0m in \u001b[0;36mSolveStokes\u001b[0;34m(X)\u001b[0m\n\u001b[1;32m     12\u001b[0m     \u001b[0mres\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mgfu\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvec\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mCreateVector\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     13\u001b[0m     \u001b[0mres\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m-\u001b[0m\u001b[0ma\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmat\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mgfu\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvec\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 14\u001b[0;31m     \u001b[0minv\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0ma\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmat\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mInverse\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfreedofs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mX\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mFreeDofs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minverse\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"umfpack\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     15\u001b[0m     \u001b[0mgfu\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvec\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdata\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0minv\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mres\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     16\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: UmfpackInverse: Numeric factorization failed."
     ]
    }
   ],
   "source": [
    "V = H1(mesh, order=2, dirichlet=\"wall|inlet|cyl\")\n",
    "Q = L2(mesh, order=1)\n",
    "print (\"V.ndof =\", V.ndof, \", Q.ndof =\", Q.ndof)\n",
    "X = FESpace([V,V,Q])\n",
    "\n",
    "gfu = SolveStokes(X)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$P^{2,+} \\times P^{1,dc}$ elements:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "V.ndof = 2496 , Q.ndof = 2382\n"
     ]
    }
   ],
   "source": [
    "V = H1(mesh, order=2, dirichlet=\"wall|inlet|cyl\")\n",
    "V.SetOrder(TRIG,3)\n",
    "V.Update()\n",
    "Q = L2(mesh, order=1)\n",
    "X = FESpace([V,V,Q])\n",
    "print (\"V.ndof =\", V.ndof, \", Q.ndof =\", Q.ndof)\n",
    "\n",
    "gfu = SolveStokes(X)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "the mini element:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "V = H1(mesh, order=1, dirichlet=\"wall|inlet|cyl\")\n",
    "V.SetOrder(TRIG,3)\n",
    "V.Update()\n",
    "Q = H1(mesh, order=1)\n",
    "X = FESpace([V,V,Q])\n",
    "\n",
    "gfu = SolveStokes(X)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## VectorH1 \n",
    "\n",
    "A vector-valued $H^1$-space: Less to type and more possibilities to explore structure and optimize. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "V = VectorH1(mesh, order=2, dirichlet=\"wall|inlet|cyl\")\n",
    "V.SetOrder(TRIG,3)\n",
    "V.Update()\n",
    "Q = L2(mesh, order=1)\n",
    "X = FESpace([V,Q])\n",
    "\n",
    "u,p = X.TrialFunction()\n",
    "v,q = X.TestFunction()\n",
    "\n",
    "a = BilinearForm(X)\n",
    "a += SymbolicBFI(InnerProduct(grad(u),grad(v))+div(u)*q+div(v)*p)\n",
    "a.Assemble()\n",
    "\n",
    "gfu = GridFunction(X)\n",
    "uin = CoefficientFunction( (1.5*4*y*(0.41-y)/(0.41*0.41), 0) )\n",
    "gfu.components[0].Set(uin, definedon=mesh.Boundaries(\"inlet\"))\n",
    "\n",
    "res = gfu.vec.CreateVector()\n",
    "res.data = -a.mat * gfu.vec\n",
    "inv = a.mat.Inverse(freedofs=X.FreeDofs(), inverse=\"umfpack\")\n",
    "gfu.vec.data += inv * res\n",
    "Draw(gfu.components[0], mesh, \"vel\")\n",
    "Draw(Norm(gfu.components[0]), mesh, \"|vel|\")\n",
    "SetVisualization(max=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stokes as a block-system\n",
    "We can now define separate bilinear-form and matrices for A and B, and combine them to a block-system:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "V = VectorH1(mesh, order=3, dirichlet=\"wall|inlet|cyl\")\n",
    "Q = H1(mesh, order=2)\n",
    "\n",
    "u,v = V.TnT()\n",
    "p,q = Q.TnT()\n",
    "\n",
    "a = BilinearForm(V)\n",
    "a += SymbolicBFI(InnerProduct(grad(u),grad(v)))\n",
    "\n",
    "b = BilinearForm(trialspace=V, testspace=Q)\n",
    "b += SymbolicBFI(div(u)*q)\n",
    "\n",
    "a.Assemble()\n",
    "b.Assemble()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Needed as preconditioner for the pressure:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "mp = BilinearForm(Q)\n",
    "mp += SymbolicBFI(p*q)\n",
    "mp.Assemble()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Two right hand sides for the two spaces:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = LinearForm(V)\n",
    "f += SymbolicLFI( CoefficientFunction( (0,x-0.5)) * v)\n",
    "f.Assemble()\n",
    "\n",
    "g = LinearForm(Q)\n",
    "g.Assemble()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Two `GridFunction`s for velocity and pressure:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "gfu = GridFunction(V, name=\"u\")\n",
    "gfp = GridFunction(Q, name=\"p\")\n",
    "uin = CoefficientFunction( (1.5*4*y*(0.41-y)/(0.41*0.41), 0) )\n",
    "gfu.Set(uin, definedon=mesh.Boundaries(\"inlet\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Combine everything to a block-system.\n",
    "`BlockMatrix` and `BlockVector` store references to the original matrices and vectors, no new large matrices are allocated. The same for the transpose matrix `b.mat.T`. It stores a wrapper for the original matrix, and replaces the call of the `Mult` function by `MultTrans`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "it =  0  err =  4.577250004746686\n",
      "it =  1  err =  2.3137027107825414\n",
      "it =  2  err =  1.9474651703872594\n",
      "it =  3  err =  1.5362811821422817\n",
      "it =  4  err =  1.490898743198753\n",
      "it =  5  err =  1.277851941805208\n",
      "it =  6  err =  1.236894318332737\n",
      "it =  7  err =  1.0832164597336917\n",
      "it =  8  err =  0.9954867877973922\n",
      "it =  9  err =  0.8694764391752701\n",
      "it =  10  err =  0.8164294710037596\n",
      "it =  11  err =  0.7270898498126048\n",
      "it =  12  err =  0.7062300357748325\n",
      "it =  13  err =  0.6435150710145074\n",
      "it =  14  err =  0.6260336874381139\n",
      "it =  15  err =  0.591910771493464\n",
      "it =  16  err =  0.577362791218001\n",
      "it =  17  err =  0.5490895491963667\n",
      "it =  18  err =  0.5299584240360132\n",
      "it =  19  err =  0.504587770754635\n",
      "it =  20  err =  0.4945747930809748\n",
      "it =  21  err =  0.478452427389186\n",
      "it =  22  err =  0.46049192930432753\n",
      "it =  23  err =  0.4367980978458769\n",
      "it =  24  err =  0.41255367261726367\n",
      "it =  25  err =  0.38040899835370073\n",
      "it =  26  err =  0.3449005478686267\n",
      "it =  27  err =  0.2985421196154886\n",
      "it =  28  err =  0.2527173796456808\n",
      "it =  29  err =  0.17951963386513986\n",
      "it =  30  err =  0.14701867167188856\n",
      "it =  31  err =  0.10086316225427124\n",
      "it =  32  err =  0.08594125971647408\n",
      "it =  33  err =  0.049535710323364915\n",
      "it =  34  err =  0.04880618304699326\n",
      "it =  35  err =  0.025962616020410965\n",
      "it =  36  err =  0.025723086935984515\n",
      "it =  37  err =  0.01628165309575614\n",
      "it =  38  err =  0.0162502272287645\n",
      "it =  39  err =  0.011879675741016503\n",
      "it =  40  err =  0.01145912264948227\n",
      "it =  41  err =  0.0077088148705381325\n",
      "it =  42  err =  0.007669232876181972\n",
      "it =  43  err =  0.0047624380822534566\n",
      "it =  44  err =  0.0047402737522261665\n",
      "it =  45  err =  0.0024617452306963774\n",
      "it =  46  err =  0.002444954453502464\n",
      "it =  47  err =  0.0015421225878978428\n",
      "it =  48  err =  0.0014860456371280366\n",
      "it =  49  err =  0.0010744692033775487\n",
      "it =  50  err =  0.0010244043188852743\n",
      "it =  51  err =  0.0006385557996683751\n",
      "it =  52  err =  0.0006291811667503349\n",
      "it =  53  err =  0.0002990274926769061\n",
      "it =  54  err =  0.0002988061757828301\n",
      "it =  55  err =  0.00011824139621187269\n",
      "it =  56  err =  0.00011462936306193394\n",
      "it =  57  err =  5.755195930067076e-05\n",
      "it =  58  err =  5.56951683872658e-05\n",
      "it =  59  err =  2.4705121755237746e-05\n",
      "it =  60  err =  2.470480085269251e-05\n",
      "it =  61  err =  1.6824109227260327e-05\n",
      "it =  62  err =  1.6572422543775576e-05\n",
      "it =  63  err =  1.316326692504858e-05\n",
      "it =  64  err =  1.2897753939698155e-05\n",
      "it =  65  err =  9.293568917613104e-06\n",
      "it =  66  err =  9.15575741204158e-06\n",
      "it =  67  err =  4.022673678037305e-06\n",
      "it =  68  err =  4.021710787407399e-06\n",
      "it =  69  err =  1.3784244536605242e-06\n",
      "it =  70  err =  1.377553514553071e-06\n",
      "it =  71  err =  6.003552147432753e-07\n",
      "it =  72  err =  5.882429968010617e-07\n",
      "it =  73  err =  2.4024447119496915e-07\n",
      "it =  74  err =  2.3576449515661013e-07\n",
      "it =  75  err =  1.0050977490275618e-07\n",
      "it =  76  err =  9.69663021570943e-08\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "basevector"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "K = BlockMatrix( [ [a.mat, b.mat.T], [b.mat, None] ] )\n",
    "C = BlockMatrix( [ [a.mat.Inverse(V.FreeDofs()), None], [None, mp.mat.Inverse()] ] )\n",
    "\n",
    "rhs = BlockVector ( [f.vec, g.vec] )\n",
    "sol = BlockVector( [gfu.vec, gfp.vec] )\n",
    "\n",
    "solvers.MinRes (mat=K, pre=C, rhs=rhs, sol=sol, initialize=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "Draw (gfu)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  },
  "nbsphinx": {
   "allow_errors": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
